import numpy as np
from PIL import Image
import tensorflow as tf
import os

train_path='./mnist_label/mnist_train/'
train_txt = './mnist_label/mnist_tarin.txt'
x_train_savepath = './mnist_label/mnist_x_train.npy'
y_train_savepath = './mnist_label/mnist_y_train.npy'

test_path = './mnist_label/mnist_test/'
test_txt = './mnist_label/mnist_test.txt'
x_test_savepath= './mnist_label/mnist_x_test.npy'
y_test_savepath= './mnist_label/mnist_y_test.npy'

def generateds(path, txt):
    f = open(txt, 'r')
    contents = f.readlines()
    f.close()
    x, y_ = [], []
    for content in contents:
        value = content.split()
        img_path = path + value[0]
        img = Image.open(img_path)
        img = np.array(img.convert('L'))
        img = img / 255.0
        x.append(img)
        y_.append(value[1])
        print('loaded:', content)
    x = np.array(x)
    y_ = np.array(y_)
    y_ = y_.astype(np.int64)
    return x, y_

if os.path.exists(x_train_savepath) and os.path.exists(y_train_savepath) and os.path.exists(
        x_test_savepath) and os.path.exists(y_test_savepath):
    print('————————————————————————————load datesets—————————————————————————————')
    x_train_save = np.load(x_train_savepath)
    y_train = np.load(y_train_savepath)
    x_test_save = np.load(x_test_savepath)
    y_test = np.load(y_test_savepath)
    x_train = np.reshape(x_train_save,(len(x_train_save),28,28))
    x_test = np.reshape(x_train_save,(len(x_test_save),28,28))
else:
    print('————————————————————————————generate datesets—————————————————————————————')
    x_train,y_train = generateds(train_path,train_txt)
    x_test, y_test = generateds(test_path, test_txt)

    print('————————————————————————————save datesets—————————————————————————————')
    x_train_save = np.reshape(x_train,(len(x_train), -1))
    x_test_save = np.reshape(x_test,(len(x_test), -1))
    np.save(x_train_savepath,x_train_save)
    np.save(y_train_savepath,y_train)
    np.save(x_test_savepath, x_test_save)
    np.save(y_train_savepath, y_test)

